import bpy
import math
import os
import cv2
from mathutils import Vector


def extract_views_from_panorama(panorama_path, output_dir, output_format, num_horizontal=4, num_vertical=3, use_directional_cameras=False, frame_number=0):
    """
    Extract multiple views from a 360 equirectangular image or video.
    
    Args:
        panorama_path (str): Path to the equirectangular image or video
        output_dir (str): Directory to save extracted views
        output_format (str): Output format ('PNG', 'JPEG', 'TIFF')
        num_horizontal (int): Number of horizontal views
        num_vertical (int): Number of vertical views
        use_directional_cameras (bool): Whether to use 12 predefined directional cameras
        frame_number (int): Frame number to extract from video (default: 0)
    
    Returns:
        list: Paths to the extracted view images
    """
    # Load the equirectangular image
    try:
        # Load image as environment texture
        bpy.ops.image.open(filepath=panorama_path)
        image_name = os.path.basename(panorama_path)
        image = bpy.data.images[image_name]
        
        # Create world environment
        if "EquirectangularWorld" not in bpy.data.worlds:
            world = bpy.data.worlds.new("EquirectangularWorld")
        else:
            world = bpy.data.worlds["EquirectangularWorld"]
        
        # Set world properties
        world.use_nodes = True
        nodes = world.node_tree.nodes
        links = world.node_tree.links
        
        # Clear existing nodes
        nodes.clear()
        
        # Create new nodes
        output_node = nodes.new(type='ShaderNodeOutputWorld')
        background_node = nodes.new(type='ShaderNodeBackground')
        env_node = nodes.new(type='ShaderNodeTexEnvironment')
        
        # Set image to environment node
        env_node.image = image
        
        # Connect nodes
        links.new(env_node.outputs['Color'], background_node.inputs['Color'])
        links.new(background_node.outputs['Background'], output_node.inputs['Surface'])
        
        # Set world to scene
        bpy.context.scene.world = world
        
    except Exception as e:
        raise Exception(f"Failed to load panorama image: {str(e)}")
    
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
    
    # List to store paths of extracted views
    extracted_views = []
    
    if use_directional_cameras:
        # Create or get the 12 directional cameras
        cameras = create_cameras_around_scene()
        
        # Set render settings
        bpy.context.scene.render.resolution_x = 1920
        bpy.context.scene.render.resolution_y = 1080
        bpy.context.scene.render.resolution_percentage = 100
        
        # Set output format
        if output_format.upper() == 'PNG':
            bpy.context.scene.render.image_settings.file_format = 'PNG'
        elif output_format.upper() == 'JPEG':
            bpy.context.scene.render.image_settings.file_format = 'JPEG'
        elif output_format.upper() == 'TIFF':
            bpy.context.scene.render.image_settings.file_format = 'TIFF'
        
        # Render from each camera
        for i, camera_obj in enumerate(cameras):
            # Set this camera as active
            bpy.context.scene.camera = camera_obj
            
            # Create file path
            filename = f"view_{frame_number:04d}_{i:02d}.{output_format.lower()}"

            filepath = os.path.join(output_dir, filename)
            
            # Set render output path
            bpy.context.scene.render.filepath = filepath
            
            # Render the view
            bpy.ops.render.render(write_still=True)
            
            # Add to list
            extracted_views.append(filepath)
    else:
        # Render views
        for h in range(num_horizontal):
            for v in range(num_vertical):
                # Calculate view direction
                # longitude: 0 to 2π (0° to 360°)
                longitude = (h / num_horizontal) * 2 * math.pi
                # latitude: -π/2 to π/2 (-90° to 90°)
                latitude = (v / num_vertical) * math.pi - (math.pi / 2)
                
                # Set up camera for this view
                camera_obj = setup_camera_for_equirectangular_view(longitude, latitude)
                
                # Set render settings
                bpy.context.scene.render.resolution_x = 1920
                bpy.context.scene.render.resolution_y = 1080
                bpy.context.scene.render.resolution_percentage = 100
                
                # Set output format
                if output_format.upper() == 'PNG':
                    bpy.context.scene.render.image_settings.file_format = 'PNG'
                elif output_format.upper() == 'JPEG':
                    bpy.context.scene.render.image_settings.file_format = 'JPEG'
                elif output_format.upper() == 'TIFF':
                    bpy.context.scene.render.image_settings.file_format = 'TIFF'
                
                # Create file path
                filename = f"view_{frame_number:04d}_h{h:02d}_v{v:02d}.{output_format.lower()}"
                filepath = os.path.join(output_dir, filename)
                
                # Set render output path
                bpy.context.scene.render.filepath = filepath
                
                # Render the view
                bpy.ops.render.render(write_still=True)
                
                # Add to list
                extracted_views.append(filepath)
    
    return extracted_views


def setup_camera_for_equirectangular_view(longitude, latitude, fov=90):
    """
    Set up a camera to render a perspective view from equirectangular coordinates.
    
    Args:
        longitude (float): Horizontal angle in radians
        latitude (float): Vertical angle in radians
        fov (float): Field of view in degrees
    """
    # Get or create camera
    if "EquirectangularCamera" in bpy.data.objects:
        camera_obj = bpy.data.objects["EquirectangularCamera"]
        camera = camera_obj.data
    else:
        camera = bpy.data.cameras.new("EquirectangularCamera")
        camera_obj = bpy.data.objects.new("EquirectangularCamera", camera)
        # Link camera to scene
        bpy.context.scene.collection.objects.link(camera_obj)
    
    # Position camera at origin
    camera_obj.location = (0, 0, 0)
    
    # Calculate rotation from longitude and latitude
    # In Blender's coordinate system:
    # X axis: forward/backward
    # Y axis: left/right
    # Z axis: up/down
    # longitude = 0 faces positive X, rotates around Z axis
    # latitude = 0 is horizontal, positive rotates upward
    
    # Convert angles to Euler rotation
    # Note: Blender uses radians for rotation
    camera_obj.rotation_euler = (
        math.pi/2 - latitude,  # X rotation (tilt up/down)
        0,                     # Y rotation (roll)
        longitude              # Z rotation (pan left/right)
    )
    
    # Set field of view
    camera.lens_unit = 'FOV'
    camera.angle = math.radians(fov)
    
    # Set camera as active
    bpy.context.scene.camera = camera_obj
    
    return camera_obj


def convert_equirectangular_to_cubemap(panorama_path, output_dir):
    """
    Convert an equirectangular image to a cubemap.
    
    Args:
        panorama_path (str): Path to the equirectangular image
        output_dir (str): Directory to save cubemap faces
    
    Returns:
        list: Paths to the cubemap face images
    """
    # Check if file is a video
    is_video = panorama_path.lower().endswith(('.mp4', '.avi', '.mov', '.mkv'))
    
    if is_video:
        # Extract specified frame from video
        temp_dir = os.path.join(output_dir, "temp_frames")
        actual_path = extract_frame_from_video(panorama_path, temp_dir, frame_number)
    else:
        # Use original image
        actual_path = panorama_path
    
    # Load the equirectangular image as environment
    try:
        # Load image as environment texture
        bpy.ops.image.open(filepath=actual_path)
        image_name = os.path.basename(actual_path)
        image = bpy.data.images[image_name]
        
        # Create world environment
        if "EquirectangularWorld" not in bpy.data.worlds:
            world = bpy.data.worlds.new("EquirectangularWorld")
        else:
            world = bpy.data.worlds["EquirectangularWorld"]
        
        # Set world properties
        world.use_nodes = True
        nodes = world.node_tree.nodes
        links = world.node_tree.links
        
        # Clear existing nodes
        nodes.clear()
        
        # Create new nodes
        output_node = nodes.new(type='ShaderNodeOutputWorld')
        background_node = nodes.new(type='ShaderNodeBackground')
        env_node = nodes.new(type='ShaderNodeTexEnvironment')
        
        # Set image to environment node
        env_node.image = image
        
        # Connect nodes
        links.new(env_node.outputs['Color'], background_node.inputs['Color'])
        links.new(background_node.outputs['Background'], output_node.inputs['Surface'])
        
        # Set world to scene
        bpy.context.scene.world = world
        
    except Exception as e:
        raise Exception(f"Failed to load panorama image: {str(e)}")
    
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
    
    # Define cubemap faces with their orientations
    # Format: (face_name, longitude, latitude)
    cubemap_orientations = [
        ('front', 0, 0),           # Front: 0° longitude, 0° latitude
        ('back', math.pi, 0),      # Back: 180° longitude, 0° latitude
        ('left', math.pi/2, 0),    # Left: 90° longitude, 0° latitude
        ('right', -math.pi/2, 0),  # Right: -90° longitude, 0° latitude
        ('top', 0, math.pi/2),     # Top: 0° longitude, 90° latitude
        ('bottom', 0, -math.pi/2)  # Bottom: 0° longitude, -90° latitude
    ]
    
    cubemap_faces = []
    
    # Render each face
    for face_name, longitude, latitude in cubemap_orientations:
        # Set up camera for this face
        camera_obj = setup_camera_for_equirectangular_view(longitude, latitude, fov=90)
        
        # Set render settings
        bpy.context.scene.render.resolution_x = 1024
        bpy.context.scene.render.resolution_y = 1024
        bpy.context.scene.render.resolution_percentage = 100
        bpy.context.scene.render.image_settings.file_format = 'PNG'
        
        # Create file path
        filename = f"cubemap_{face_name}.png"
        filepath = os.path.join(output_dir, filename)
        
        # Set render output path
        bpy.context.scene.render.filepath = filepath
        
        # Render the face
        bpy.ops.render.render(write_still=True)
        
        # Add to list
        cubemap_faces.append(filepath)
    
    return cubemap_faces


def extract_frame_from_video(video_path, output_dir, frame_number=0):
    """
    Extract a specific frame from a video file.
    
    Args:
        video_path (str): Path to the video file
        output_dir (str): Directory to save the extracted frame
        frame_number (int): Which frame to extract (default: 0 for first frame)
    
    Returns:
        str: Path to the extracted frame image
    """
    try:
        # Create output directory if it doesn't exist
        os.makedirs(output_dir, exist_ok=True)
        
        # Open video file
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise Exception("Could not open video file")
        
        # Get total frame count
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        # Ensure frame_number is within bounds
        if frame_number >= total_frames:
            frame_number = total_frames - 1
        
        # Set frame position and read frame
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
        ret, frame = cap.read()
        
        if not ret:
            raise Exception("Could not read frame from video")
        
        # Save frame as image
        frame_filename = f"frame_{frame_number:04d}.jpg"
        frame_path = os.path.join(output_dir, frame_filename)
        cv2.imwrite(frame_path, frame)
        
        cap.release()
        return frame_path
        
    except Exception as e:
        raise Exception(f"Failed to extract frame from video: {str(e)}")


def extract_views_from_video_frames(video_path, output_dir, output_format, frame_interval=1, use_directional_cameras=True):
    """
    Extract views from multiple frames of a video file sequentially.
    
    Args:
        video_path (str): Path to the video file
        output_dir (str): Directory to save extracted views
        output_format (str): Output format ('PNG', 'JPEG', 'TIFF')
        frame_interval (int): Process every Nth frame (1 = all frames)
        use_directional_cameras (bool): Whether to use 12 predefined directional cameras
    
    Returns:
        dict: Dictionary mapping frame numbers to lists of extracted view paths
    """
    try:
        # Open video file
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise Exception("Could not open video file")
        
        # Get video properties
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        fps = cap.get(cv2.CAP_PROP_FPS)
        
        cap.release()
        
        # Create main output directory
        os.makedirs(output_dir, exist_ok=True)
        
        # Dictionary to store results
        frame_results = {}
        
        # Process frames at intervals
        for frame_num in range(0, total_frames, frame_interval):
            print(f"Processing frame {frame_num}/{total_frames}")
            
            # Create subdirectory for this frame
            frame_dir = os.path.join(output_dir, f"frames")
            os.makedirs(frame_dir, exist_ok=True)
            
            # Extract frame and process views
            try:
                # Create temporary directory for frame extraction
                temp_dir = os.path.join(output_dir, "temp_frames")
                frame_path = extract_frame_from_video(video_path, temp_dir, frame_num)
                
                # Extract views from this frame
                views = extract_views_from_panorama(
                    frame_path,
                    frame_dir,
                    output_format,
                    use_directional_cameras=use_directional_cameras,
                    frame_number=frame_num
                )
                
                frame_results[frame_num] = views
                
            except Exception as e:
                print(f"Failed to process frame {frame_num}: {str(e)}")
                continue
        
        return frame_results
        
    except Exception as e:
        raise Exception(f"Failed to extract views from video frames: {str(e)}")

def set_panorama_as_environment(panorama_path):
    """
    Set the panorama image or video as the world environment texture.
    
    Args:
        panorama_path (str): Path to the equirectangular image or video
    """
    try:
        # Check if file is a video
        is_video = panorama_path.lower().endswith(('.mp4', '.avi', '.mov', '.mkv'))
        
        if is_video:
            # Extract first frame from video
            temp_dir = os.path.join(os.path.dirname(panorama_path), "temp_frames")
            frame_path = extract_frame_from_video(panorama_path, temp_dir)
            
            # Use the extracted frame as environment
            actual_path = frame_path
        else:
            # Use original image
            actual_path = panorama_path
        
        # Load the equirectangular image
        bpy.ops.image.open(filepath=actual_path)
        image_name = os.path.basename(actual_path)
        image = bpy.data.images[image_name]
        
        # Create world environment
        if "EquirectangularWorld" not in bpy.data.worlds:
            world = bpy.data.worlds.new("EquirectangularWorld")
        else:
            world = bpy.data.worlds["EquirectangularWorld"]
        
        # Set world properties
        world.use_nodes = True
        nodes = world.node_tree.nodes
        links = world.node_tree.links
        
        # Clear existing nodes
        nodes.clear()
        
        # Create new nodes
        output_node = nodes.new(type='ShaderNodeOutputWorld')
        background_node = nodes.new(type='ShaderNodeBackground')
        env_node = nodes.new(type='ShaderNodeTexEnvironment')
        
        # Set image to environment node
        env_node.image = image
        env_node.projection = 'EQUIRECTANGULAR'
        
        # Connect nodes
        links.new(env_node.outputs['Color'], background_node.inputs['Color'])
        links.new(background_node.outputs['Background'], output_node.inputs['Surface'])
        
        # Set world to scene
        bpy.context.scene.world = world
        
        # Update viewport shading to show environment
        for area in bpy.context.screen.areas:
            if area.type == 'VIEW_3D':
                for space in area.spaces:
                    if space.type == 'VIEW_3D':
                        space.shading.type = 'MATERIAL'
                        space.shading.use_scene_lights = True
                        space.shading.use_scene_world = True
                        break
        
        return True
        
    except Exception as e:
        raise Exception(f"Failed to set panorama as environment: {str(e)}")


def create_cameras_around_scene():
    """
    Create 12 cameras positioned around the scene, organized in 3 groups:
    - Group 1: 4 cameras facing upward at 45 degrees
    - Group 2: 4 cameras facing horizontally
    - Group 3: 4 cameras facing downward at 45 degrees
    """
    # Remove existing directional cameras if any
    for i in range(18):
        cam_name = f"DirectionalCamera_{i:02d}"
        if cam_name in bpy.data.objects:
            bpy.data.objects.remove(bpy.data.objects[cam_name], do_unlink=True)
    
    # Define 12 directions organized in 3 groups of 4 cameras each
    # Each group has cameras at 90-degree intervals (0°, 90°, 180°, 270°)
    directions = [
        # Group 1: Upward 45 degrees (x_rot = π/4)
        (0, math.pi/4,  0),              # Front-Up
        (math.pi/3, math.pi/4, 0),       # Right-Up
        (math.pi*2/3, math.pi/4, 0),      # Right-Up
        (math.pi, math.pi/4, 0),        # Back-Up
        (math.pi*4/3, math.pi/4, 0),    # Left-Up
        (math.pi*5/3, math.pi/4, 0),    # Left-Up

        # # Group 2: Horizontal (x_rot = 0)
        (0, math.pi/2, 0),                      # Front
        (math.pi/3, math.pi/2, 0),              # Right
        (math.pi*2/3, math.pi/2, 0),            # Back
        (math.pi,math.pi/2, 0),                 # Left
        (math.pi*4/3, math.pi/2, 0),            # Left
        (math.pi*5/3, math.pi/2, 0),            # Left
        
        # # Group 3: Downward 45 degrees (x_rot = -π/4)
        (0, 3 * math.pi/4, 0),             # Front-Down
        (math.pi/3, 3 * math.pi/4, 0),     # Right-Down
        (math.pi*2/3, 3 * math.pi/4, 0),       # Back-Down
        (math.pi, 3 * math.pi/4, 0),       # Back-Down
        (math.pi*4/3, 3 * math.pi/4, 0),       # Left-Down
        (math.pi*5/3, 3 * math.pi/4, 0),       # Left-Down

    ]
    
    cameras = []
    
    for i, (z_rot, x_rot, y_rot) in enumerate(directions):
        # Create camera data
        cam_data = bpy.data.cameras.new(f"DirectionalCamera_{i:02d}")
        cam_data.lens_unit = 'FOV'
        cam_data.angle = math.radians(120)  # 120° field of view
        
        # Create camera object
        cam_obj = bpy.data.objects.new(f"DirectionalCamera_{i:02d}", cam_data)
        
        # Position camera at origin (0, 0, 0)
        cam_obj.location = (0, 0, 0)
        
        # Set rotation (Euler angles in radians)
        cam_obj.rotation_euler = (x_rot, y_rot, z_rot)
        
        # Link camera to scene
        bpy.context.scene.collection.objects.link(cam_obj)
        
        cameras.append(cam_obj)
    
    # Set the first camera as active
    if cameras:
        bpy.context.scene.camera = cameras[0]
    
    return cameras